// modified version of the JasminVisitor class included in the BCEL project
// see http://jakarta.apache.org/bcel/ for more details

// the following line has been added for integration in the tinapoc project
package tina.utils;

import org.apache.bcel.classfile.Deprecated;
import org.apache.bcel.classfile.*;
import org.apache.bcel.generic.*;
import java.io.*;
import java.util.*;
import org.apache.bcel.Constants;
import org.apache.bcel.Repository;


/**
 * Disassemble Java class object into the <A HREF="http://www.cat.nyu.edu/meyer/jasmin">
 * JASMIN</A> format.
 * 
 * modified by D. Reynaud for the tinapoc project (http://tinapoc.sourceforge.net)
 *
 * @version $Id$
 * @author <A HREF="mailto:m.dahm@gmx.de">M. Dahm</A>
 */
public class NeoJasminVisitor extends org.apache.bcel.classfile.EmptyVisitor {
    private JavaClass       clazz;
    private PrintWriter     out;
    private String          class_name;
    private ConstantPoolGen cp;
    private boolean         oldjasmin,
                            nolabels,
                            printoffsets;
    private Attribute[] attr;

    public NeoJasminVisitor(JavaClass clazz, OutputStream out,
             boolean oldjasmin, boolean nolabels, boolean printoffsets) {
        this.clazz = clazz;
        this.attr = clazz.getAttributes();
        this.out = new PrintWriter(out);
        this.oldjasmin = oldjasmin;
        this.nolabels = nolabels;
        this.printoffsets = printoffsets;
        class_name = clazz.getClassName();
        cp = new ConstantPoolGen(clazz.getConstantPool());
    }

    /**
     * Start traversal using DefaultVisitor pattern.
     */
    public void disassemble() {
        new org.apache.bcel.classfile.DescendingVisitor(clazz, this).visit();
        out.close();
    }

    public void visitJavaClass(JavaClass clazz) {
        out.println("; Produced by NeoJasminVisitor (tinapoc)");
        out.println("; http://tinapoc.sourceforge.net");
        out.println("; The original JasminVisitor is part of the BCEL");
        out.println("; http://jakarta.apache.org/bcel/");
        out.println("; " + new Date() + "\n");

        int flags = clazz.getAccessFlags();
        if(!oldjasmin)
            out.println(".bytecode " + clazz.getMajor() + "." + clazz.getMinor());
        out.println(".source " + clazz.getSourceFileName());
        out.println("." + Utility.classOrInterface(flags) + " " +
            Utility.accessToString(flags, true) + " " +
            ((isEnum(flags) && !oldjasmin) ? "enum " : "") +
            ((isAnnotation(flags) && !oldjasmin)? "annotation " : "") +
            clazz.getClassName().replace('.', '/'));
        out.println(".super " + clazz.getSuperclassName().replace('.', '/'));

        String signature = getClassSignature();
        if(signature!=null && !oldjasmin)
            out.println(".signature \""+signature+"\"");

        String debug = getSourceDebugExtension();
        if(debug!=null && !oldjasmin)
            out.println(".debug \""+
                StringNormalizer.removeSpecialCharacters(debug)+"\"");

        String enc = getEnclosingMethod();
        if(enc!=null && !oldjasmin)
            out.println(".enclosing method \""+enc+"\"");

        String[] interfaces = clazz.getInterfaceNames();

        for(int i=0; i < interfaces.length; i++)
            out.println(".implements " + interfaces[i].replace('.', '/'));

        out.print("\n");
    }

    public void visitField(Field field) {
        int flags = field.getAccessFlags();
        String sig = getSignature(field.getAttributes());
        out.println(".field " + Utility.accessToString(flags) + " " +
                  ((isEnum(flags) && !oldjasmin) ? "enum " : "") +
                  field.getName() + " " + field.getSignature() +
                  ((sig!=null && !oldjasmin) ? " signature \""+sig+"\"" : ""));
    }

    public void visitConstantValue(ConstantValue cv) {
        out.println(" = " + cv);
    }

    private Method _method;

    /**
     * Unfortunately Jasmin expects ".end method" after each method. Thus we've to check
     * for every of the method's attributes if it's the last one and print ".end method"
     * then.
     */
    private final void printEndMethod(Attribute attr) {
        Attribute[] attributes = _method.getAttributes();

        if(attr == attributes[attributes.length - 1])
            out.println(".end method");
    }

    public void visitDeprecated(Deprecated attribute) { printEndMethod(attribute); }

    public void visitSynthetic(Synthetic attribute) {
        if(_method != null)
            printEndMethod(attribute);
    }

    public void visitMethod(Method method) {
        this._method = method; // Remember for use in subsequent visitXXX calls

        int flags = _method.getAccessFlags();
        out.println("\n\n\n.method " + Utility.accessToString(flags).replace("volatile",
            oldjasmin?"":"bridge").replace("transient", oldjasmin?"":"varargs")
            + " " + _method.getName() + _method.getSignature());

        Attribute[] attributes = _method.getAttributes();
        String sig;
        if((attributes == null) || (attributes.length == 0))
            out.println(".end method");
        else if((sig=getSignature(attributes))!=null && !oldjasmin)
            out.println("    .signature \""+sig+"\"");
    }

    public void visitExceptionTable(ExceptionTable e) {
        String[] names = e.getExceptionNames();
        for(int i=0; i < names.length; i++)
            out.println("    .throws " + names[i].replace('.', '/'));

        printEndMethod(e);
    }

    private Hashtable<InstructionHandle,String> map;

    public void visitCode(Code code) {
        int label_counter = 0;

        out.println("    .limit stack " + code.getMaxStack());
        out.println("    .limit locals " + code.getMaxLocals());

        MethodGen           mg  = new MethodGen(_method, class_name, cp);
        InstructionList     il  = mg.getInstructionList();
        InstructionHandle[] ihs = il.getInstructionHandles();

        /* Pass 1: Give all referenced instruction handles a symbolic name, i.e. a
         * label.
         */
        map = new Hashtable<InstructionHandle,String>();

        for(int i=0; i < ihs.length; i++) {
            if(ihs[i] instanceof BranchHandle) {
                BranchInstruction bi = (BranchInstruction)ihs[i].getInstruction();
  
                if(bi instanceof Select) { // Special cases LOOKUPSWITCH and TABLESWITCH
                    InstructionHandle[] targets = ((Select)bi).getTargets();
        
                    for(int j=0; j < targets.length; j++)
                        put(targets[j], "Label" + label_counter++ + ":");
                }

                InstructionHandle ih = bi.getTarget();
                put(ih, "Label" + label_counter++ + ":");
            }
        }

        LocalVariableGen[] lvs = mg.getLocalVariables();
        for(int i=0; i < lvs.length; i++) {
            InstructionHandle ih = lvs[i].getStart();
            put(ih, "Label" + label_counter++ + ":");
            ih = lvs[i].getEnd();
            put(ih, "Label" + label_counter++ + ":");    
        }
        
        CodeExceptionGen[] ehs = mg.getExceptionHandlers();
        for(int i=0; i < ehs.length; i++) {
            CodeExceptionGen  c  = ehs[i];
            InstructionHandle ih = c.getStartPC();

            put(ih, "Label" + label_counter++ + ":");
            ih = c.getEndPC().getNext();   // getEndPc() is INCLUSIVE !!!!
            put(ih, "Label" + label_counter++ + ":");    
            ih = c.getHandlerPC();
            put(ih, "Label" + label_counter++ + ":");    
        }

        LineNumberGen[] lns = mg.getLineNumbers();
        for(int i=0; i < lns.length; i++) {
            InstructionHandle ih = lns[i].getInstruction();
            put(ih, ".line " + lns[i].getSourceLine());
        }
 
        /* Pass 2: Output code.
         */
        for(int i=0; i < lvs.length; i++) {
            LocalVariableGen l = lvs[i];
            out.println("    .var " + l.getIndex() + " is " + l.getName() +" "+
            l.getType().getSignature() + 
            " from " + get(l.getStart()) + 
            " to " + get(l.getEnd()));
        }

        out.print("\n    ");
        
        for(int i=0; i < ihs.length; i++) {
            InstructionHandle ih   = ihs[i];
            Instruction       inst = ih.getInstruction();
            String            str  = (String)map.get(ih);
            
            if(str != null) {
                if(!str.startsWith("Label") || !nolabels)
                    out.println("\n    "+str);
            }
            int pc = ih.getPosition();

            if(inst instanceof BranchInstruction) {
                if(inst instanceof Select) { // Special cases LOOKUPSWITCH and TABLESWITCH
                    Select s = (Select)inst;
                    int[] matchs = s.getMatchs();
                    InstructionHandle[] targets = s.getTargets();
                    int[] targetso = s.getIndices();
        
                    if(s instanceof TABLESWITCH) {
                        if(oldjasmin || !printoffsets)
                            out.printf("    tableswitch %d %d\n",
                                matchs[0],
                                matchs[matchs.length - 1]);
                        else
                            out.printf("    %4d: tableswitch %d %d\n",
                                pc,
                                matchs[0],
                                matchs[matchs.length - 1]);
                        for(int j=0; j < targets.length; j++) {
                            if(nolabels)
                                out.println("            " + targetso[j]);
                            else
                                out.println("            " + get(targets[j]));
                        }

                    } else { // LOOKUPSWITCH
                        if(oldjasmin || !printoffsets)
                            out.printf("    lookupswitch\n");
                        else
                            out.printf("    %4d: lookupswitch\n", pc);

                        for(int j=0; j < targets.length; j++) {
                            if(nolabels)
                                out.println("            " + matchs[j]
                                            + " : " + targetso[j]);
                            else
                                out.println("            " + matchs[j]
                                            + " : " + get(targets[j]));
                        }
                    }

               // Applies for both
                    out.print("            ");
                    if(nolabels)
                        out.println("default: " + s.getIndex());
                    else
                        out.println("default: " + get(s.getTarget()));

                } else {
                    BranchInstruction bi = (BranchInstruction)inst;
                    ih  = bi.getTarget();
                    str = get(ih);
                    if(nolabels) {
                        if(oldjasmin || !printoffsets)
                            out.printf("    %s %d\n",
                                Constants.OPCODE_NAMES[bi.getOpcode()],
                                bi.getIndex());
                        else
                            out.printf("    %4d: %s %d\n",
                                pc, 
                                Constants.OPCODE_NAMES[bi.getOpcode()],
                                bi.getIndex());
                    } else {
                        if(oldjasmin || !printoffsets)
                            out.printf("    %s %s\n",
                                Constants.OPCODE_NAMES[bi.getOpcode()],
                                str);
                        else
                            out.printf("    %4d: %s %s\n",
                                pc,
                                Constants.OPCODE_NAMES[bi.getOpcode()],
                                str);
                    }
                }
            } // end if instanceof BranchInstruction
            else {
                if(oldjasmin || !printoffsets)
                    out.printf("    %s\n", inst.toString(cp.getConstantPool()));
                else
                    out.printf("    %4d: %s\n",
                        pc,
                        inst.toString(cp.getConstantPool()));
            }
        } // end for
        
        out.print("\n");

        for(int i=0; i < ehs.length; i++) {
            CodeExceptionGen c = ehs[i];
            ObjectType caught = c.getCatchType();
            String class_name = (caught == null)?    // catch any exception, used when compiling finally
        "all" : caught.getClassName().replace('.', '/');

            out.println(".catch " + class_name + " from " +
            get(c.getStartPC()) + " to " + get(c.getEndPC().getNext()) +
            " using " + get(c.getHandlerPC()));
        }
        printStackMapAttribute(code);
        printEndMethod(code);
    }
 
    private final String get(InstructionHandle ih) {
        if(nolabels) {
            return String.valueOf(ih.getPosition());
        } else {
            String str = new StringTokenizer((String)map.get(ih), "\n").nextToken();
            return str.substring(0, str.length() - 1);
        }
    }

    private final void put(InstructionHandle ih, String line) {
        String str = (String)map.get(ih);

        if(str == null)
            map.put(ih, line);
        else {
            if(line.startsWith("Label") || str.endsWith(line)) // Already have a label in the map
    return;

            map.put(ih, str + "\n" + line); // append
        }
    }

    private final boolean isEnum(int flags) {
        if((flags & 0x4000)==0x4000)
            return true;
        else
            return false;
    }

    private final boolean isAnnotation(int flags) {
        if((flags & 0x2000)==0x2000)
            return true;
        else
            return false;
    }

    private final boolean isBridge(int flags) {
        if((flags & 0x0040)==0x0040)
            return true;
        else
            return false;
    }

    private final boolean isVarargs(int flags) {
        if((flags & 0x0080)==0x0080)
            return true;
        else
            return false;
    }

    private final String getClassSignature() {
        return getSignature(attr);
    }

    private final String getSignature(Attribute[] attr) {
        for(int i=0; i<attr.length; i++) {
            if(attr[i] instanceof Signature)
                return ((Signature)attr[i]).getSignature();
        }
        return null;
    }

    private final String getSourceDebugExtension() {
        for(int i=0; i<attr.length; i++) {
            if(attr[i] instanceof Unknown) {
                if(((Unknown)attr[i]).getName().equals("SourceDebugExtension")) {
                    DataInputStream in = new DataInputStream(new ByteArrayInputStream(((Unknown)attr[i]).getBytes()));
                    try {
                        return in.readUTF();
                    } catch(IOException e) {
                        e.printStackTrace();
                    }
                }
            }
        }
        return null;
    }

    private final String getEnclosingMethod() {
        for(int i=0; i<attr.length; i++) {
            if(!(attr[i] instanceof Unknown))
                continue;

            if(!((Unknown)attr[i]).getName().equals("EnclosingMethod"))
                continue;

            DataInputStream in = new DataInputStream(new ByteArrayInputStream(((Unknown)attr[i]).getBytes()));
            StringBuilder b = new StringBuilder();
            ConstantPool cpool = cp.getConstantPool();

            try {
                int idx = in.readUnsignedShort();
                Constant c = cp.getConstant(idx);
                if(!(c instanceof ConstantClass))
                    throw new RuntimeException("expecting CONSTANT_Class at index "+idx);

                b.append(((ConstantClass)c).getConstantValue(cpool));
                idx = in.readUnsignedShort();
                if(idx == 0) 
                    return b.toString();

                c = cp.getConstant(idx);
                if(!(c instanceof ConstantNameAndType))
                    throw new RuntimeException("expecting CONSTANT_NameAndType at index "+idx);

                b.append("/");
                b.append(((ConstantNameAndType)c).getName(cpool));
                b.append(((ConstantNameAndType)c).getSignature(cpool));
                return b.toString();
            } catch(IOException e) {
                e.printStackTrace();
            }
        } // for
        return null;
    }

    private void printStackMapAttribute(Code code) {
        Attribute[] a = code.getAttributes();
        StackMap stackmap = null;

        for(int i=0; i<a.length; i++) {
            if(a[i] instanceof StackMap) {
                stackmap = (StackMap)a[i];
                break;
            }
        }
        if(stackmap == null) // no stackmap found
            return;

        StackMapEntry[] map = stackmap.getStackMap();
        out.println("\n; StackMap attribute found with "+map.length+" entries");
        for(int i=0; i<map.length; i++) {
            out.println("\n    .stack");
            out.println("        offset "+map[i].getByteCodeOffset());

            StackMapType[] locals = map[i].getTypesOfLocals();
            StackMapType[] stack = map[i].getTypesOfStackItems();
            printStackMapItems(locals, true);
            printStackMapItems(stack, false);

            out.println("    .end stack");
        }
    }

    private void printStackMapItems(StackMapType[] items, boolean locals) {
        String[] tags = {"Top", "Integer", "Float", "Double", "Long", "Null",
                         "UninitializedThis", "Object", "Uninitialized"};

        for(int j=0; j<items.length; j++) {
            out.print("        ");
            out.print(locals ? "locals " : "stack ");
            out.print(tags[items[j].getType()]);
            if(items[j].hasIndex()) {
                int idx = items[j].getIndex();
                if(items[j].getType() == 8) // Uninitialized
                    out.print(" "+idx);
                else {
                    Constant c = cp.getConstant(idx);
                    if(!(c instanceof ConstantClass))
                        throw new RuntimeException("Expecting CONSTANT_Class"
                                      +" at index "+idx);

                    ConstantClass cls = (ConstantClass)c;
                    idx = cls.getNameIndex();
                    c = cp.getConstant(idx);
                    if(!(c instanceof ConstantUtf8))
                        throw new RuntimeException("Expecting CONSTANT_Utf8"
                                      +" at index "+idx);
                    out.print(" "+((ConstantUtf8)c).getBytes());
                }
            }
            out.println();
        }
    }
}
